import random
from functools import partial
import sympy as sp

class RandomGenerator:
    '''
    Implements objects that represent sets of arbitrary objects and generate 
    random elements from these sets.
    '''
    def __init__(self):
        self._names = {}
        self._created = []

    def create_new(self, *args):
        '''Generate a random element from the given set.
        
        Must be overridden in child classes'''
        raise NotImplementedError('must be implemented in subclasses')

    def new(self, name=None, *args, **kwds):
        '''
        Returns a new random element of the given set. The first argument is an 
        optional name that can be used to reference the newly created object.
        '''
        new = self.create_new(*args, **kwds)
        if name is not None:
            self._names[name] = new
        self._created.append(new)
        return new

    def __call__(self, name=None, *args, **kwds):
        '''Same as obj.new(name, *args, **kwds)'''

        return self.new(name=name, *args, **kwds)

    # Magical methods ----------------------------------------------------------
    def __str__(self):
        return str(self.new())

    def __getitem__(self, idx):
        if isinstance(idx, int):
            return self._created[idx]
        else:
            return self._names[idx]

    # Registering subclasses ---------------------------------------------------
    TEMPLATES_NS = {}

    @classmethod
    def register(cls, tt, name=None):
        '''Register a subclass as a generator object for newly created template
        namespaces'''

        name = name or tt.__name__
        cls.TEMPLATES_NS[name] = tt

    @classmethod
    def new_namespace(cls):
        '''Returns a dictionary with a new template namespace initialized from 
        all registered generator classes'''

        return {name: tt() for name, tt in cls.TEMPLATES_NS.items()}

    @classmethod
    def from_func(cls, func=None, name=None, *, clsname=None, register=True):
        '''A factory function that creates subclasses of RandomGenerator by 
        overriding the ``create_new`` method'''

        # Support a decorator call such as
        #
        # @RandomGenerator.from_func("my_name")
        # def foo(...):
        #    ...
        #
        if not callable(func):
            if name is not None:
                raise TypeError("got multiple values for argument 'name'")
            return partial(cls.from_func, name=func, register=register)

        # Create class and register
        clsname = clsname or func.__name__
        tt = type(clsname, (RandomGenerator,), {'create_new': staticmethod(func)})
        tt.__doc__ = func.__doc__

        if register:
            # respect convention that a class "XXGen" is named as "XX"
            name = name or clsname[:-3]
            cls.register(tt, name)
        return tt

#===============================================================================
# Numerical sets
#===============================================================================
@RandomGenerator.from_func
def NNGen(nonzero=True):
    '''Natural numbers generator.'''

    N = random.randint(int(nonzero), 100)
    return sp.Integer(N)

@RandomGenerator.from_func
def ZZGen(force_pos=False, *args, **kwds):
    '''Integers generator.'''

    N = NNGen.create_new(*args, **kwds)
    if random.randint(0, 1):
        N *= -1
    return N

@RandomGenerator.from_func
def RRGen(*args, **kwds):
    '''Real/float numbers generator.'''

    return sp.Float(ZZGen.create_new(*args, **kwds) / 100)

@RandomGenerator.from_func
def QQGen(*args, **kwds):
    '''Real/float numbers generator.'''
    return ZZGen.create_new(*args, **kwds) / ZZGen.create_new(nonzero=True, *args, **kwds)


if __name__ == '__main__':
    print(RandomGenerator.new_namespace())
